# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
Backbone modules.
"""
from collections import OrderedDict

import torch
import torch.nn.functional as F
import torchvision
from torch import nn
from torchvision.models._utils import IntermediateLayerGetter
from typing import Dict, List

from util.misc import NestedTensor, is_main_process

from .position_encoding import build_position_encoding


class FrozenBatchNorm2d(torch.nn.Module):
    """
    BatchNorm2d where the batch statistics and the affine仿射、映射 parameters are fixed.

    Copy-paste from torchvision.misc.ops with added eps before rqsrt,
    without which any other models than torchvision.models.resnet[18,34,50,101]
    produce nans.
    """

    def __init__(self, n):
        super(FrozenBatchNorm2d, self).__init__()  #初始化父类的属性、方法  子类__init__()会覆盖父类的__init__()
        self.register_buffer("weight", torch.ones(n))  #用于注册非参数张量，专门用于管理模型中 “需要跟着模型走，但不需要训练” 的张量的工具
        self.register_buffer("bias", torch.zeros(n))
        self.register_buffer("running_mean", torch.zeros(n))
        self.register_buffer("running_var", torch.ones(n))

    #state_dict为一字典，存储了模型的参数（nn.Parameter）和缓冲区（register_buffer 注册的张量），当加载预训练模型时，PyTorch 会默认将 state_dict 中的键值对与模型中的参数 / 缓冲区一一对应赋值。
    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
                              missing_keys, unexpected_keys, error_msgs):
        #1. 构造 BatchNorm 中 "num_batches_tracked" 缓冲区的键名
        num_batches_tracked_key = prefix + 'num_batches_tracked'
        # 2. 如果 state_dict 中存在该键，则删除它
        if num_batches_tracked_key in state_dict:
            del state_dict[num_batches_tracked_key]
        # 3. 调用父类的 _load_from_state_dict 方法，执行剩余的加载逻辑
        super(FrozenBatchNorm2d, self)._load_from_state_dict(
            state_dict, prefix, local_metadata, strict,
            missing_keys, unexpected_keys, error_msgs)

    def forward(self, x):
        # move reshapes to the beginning
        # to make it fuser-friendly
        w = self.weight.reshape(1, -1, 1, 1)
        b = self.bias.reshape(1, -1, 1, 1)
        rv = self.running_var.reshape(1, -1, 1, 1)
        rm = self.running_mean.reshape(1, -1, 1, 1)
        eps = 1e-5 #微小常数0.00001，避免数值计算中除以0或精度问题
        scale = w * (rv + eps).rsqrt() #rsqrt()方法用于计算平方根的倒数    (rv + eps).rsqrt()就是计算(rv + eps)的开根号分之一
        bias = b - rm * scale
        return x * scale + bias


class BackboneBase(nn.Module):

    def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool):
        super().__init__()
        for name, parameter in backbone.named_parameters(): 
            #控制模型backbone中的参数是否需要计算梯度
            if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
                parameter.requires_grad_(False)#冻结参数：不计算梯度，不更新
        #定义特征提取层，返回中间基层还是最后一层
        if return_interm_layers: 
            return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"}
        else:
            return_layers = {'layer4': "0"}
        self.body = IntermediateLayerGetter(backbone, return_layers=return_layers) 
        self.num_channels = num_channels

    def forward(self, tensor_list: NestedTensor):
        xs = self.body(tensor_list.tensors)#将输入图像传入 backbone 网络，提取指定层的特征。
        out: Dict[str, NestedTensor] = {} #类注释写法，创建一个名为out的变量，out是一个空字典，键为str类型，值为NestedTensor类型
        for name, x in xs.items():#从 backbone 提取的特征图生成对应的对应的掩码（mask）
            m = tensor_list.mask
            assert m is not None
            #m[None]: 给掩码添加一个维度，从 [batch_size, H, W] 变为 [1, batch_size, H, W]），因为 F.interpolate 要求输入是 4 维张量（[batch, channel, H, W]）。
            #.float()：将bool类型转成float(),interpolate()方法不支持bool类型
            #shape[-2:]：特征图的高度和宽度
            #to(torch.bool):转换成原来的bool类型
            #[0]:移除m[None]添加的维度
            mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0]
            out[name] = NestedTensor(x, mask)#是一个自定义数据结构（常见于 DETR 等检测框架），用于同时存储特征图 x 和对应的掩码 mask。方便后续网络获取特征和mask。
        return out


#Backbone网络初始化
class Backbone(BackboneBase):  
    """ResNet backbone with frozen BatchNorm."""  #冻结ResNet的BN层
    def __init__(self, name: str,
                 train_backbone: bool,
                 return_interm_layers: bool,
                 dilation: bool):
        #getattr(object, name) 是一个内置函数，用于动态获取对象的属性或方法。
        #它允许通过字符串形式的名称来访问对象的属性，而不是直接使用点语法（如 obj.attr）
        #这在需要动态处理属性时非常有用
        #object：要操作的对象（可以是类实例、类本身等）。
        #name：字符串类型，表示要获取的属性或方法的名称。
        backbone = getattr(torchvision.models, name)(
            replace_stride_with_dilation=[False, False, dilation], #表示将 ResNet 最后三个卷积块的 stride 替换为空洞卷积（仅最后一个块生效，前两个不替换），用于控制特征图尺寸
            pretrained=is_main_process(), #表示仅在主进程加载预训练权重，避免重复加载
            norm_layer=FrozenBatchNorm2d)  #使用FrozenBatchNorm2d代替原来BN层
        num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
        super().__init__(backbone, train_backbone, num_channels, return_interm_layers)


class Joiner(nn.Sequential):
    def __init__(self, backbone, position_embedding):
        super().__init__(backbone, position_embedding)

    def forward(self, tensor_list: NestedTensor):
        xs = self[0](tensor_list) #对象索引+调用可调用对象
        out: List[NestedTensor] = []
        pos = []
        for name, x in xs.items():
            out.append(x)
            # position encoding
            pos.append(self[1](x).to(x.tensors.dtype))

        return out, pos


def build_backbone(args):
    position_embedding = build_position_encoding(args)
    train_backbone = args.lr_backbone > 0
    return_interm_layers = args.masks
    backbone = Backbone(args.backbone, train_backbone, return_interm_layers, args.dilation)
    model = Joiner(backbone, position_embedding)
    model.num_channels = backbone.num_channels
    return model
